const express = require("express");
const router = express.Router();
const { spawn, spawnSync, exec } = require("child_process");
const path = require("path");
const verifyStatus = require("../middleware/verifyStatus");

router.get("/run", verifyStatus(), (req, res) => {
  const DROP_SIZE = 1000;
  const BUFFER_SIZE = 560;
  let buffer = [];
  let curSum = 0;

  const args = `--exp-name ${req.query.expName} --total-timesteps ${req.query.totalTimesteps} --learning-rate ${req.query.learningRate} --buffer-size ${req.query.bufferSize} --gamma ${req.query.gamma} --batch-size ${req.query.batchSize} --start-e ${req.query.startEpsilon} --end-e ${req.query.endEpsilon} --exploration-fraction ${req.query.explorationFraction} --learning-starts ${req.query.learningStarts}`;
  // const child = spawn("python", ["train.py"], {
  //   cwd: path.join(process.cwd(), "app"),
  // });
  const child = exec(
    `bash -c 'source venv/bin/activate && python ./CV-RL-PR/RL/DQNv2/dqn_main.py ${args}'`,
    {
      cwd: path.join(process.cwd(), "app"),
    }
  );
  child.stdout.on("data", (data) => {
    const msg = data.toString().trim();
    const strs = msg.split(" ");
    if (!strs[2]?.startsWith("global_step")) {
      return;
    }
    const nowStep = strs[2]?.split("=")[1]?.split("/")[0];
    const reward = Number.parseInt(strs[5]?.split("=")[1]);

    if (nowStep > DROP_SIZE) {
      buffer.push(reward);
      curSum += reward;
      if (buffer.length > BUFFER_SIZE) {
        curSum -= buffer.shift();
      }
    }

    const info = {
      x: strs[1],
      y: strs[5]?.split("=")[1],
      avg: buffer.length === BUFFER_SIZE ? curSum / BUFFER_SIZE : 0,
      image: strs[3]?.split("=")[1],
      cvmodel: strs[4]?.split("=")[1],
    };
    const seq = JSON.stringify(info);
    req.app.locals.wsTrainLog?.send(msg);
    req.app.locals.wsTrainChart?.send(seq);
  });

  child.stderr.on("data", (data) => {
    console.log(`${data}`);
    req.app.locals.wsTrainLog.send(`${data}`);
  });

  child.on("close", (code) => {
    console.log(`子进程退出,退出码 ${code}`);
    req.app.locals.wsTrainLog.send(`exit: ${code}`);
    req.app.locals.status = "WAITING";
  });
  req.app.locals.status = "TRAINING";
  req.app.locals.trainProcess = child;
  return res.json({ code: 0, msg: "训练任务启动成功" });
});

router.get("/stop", (req, res) => {
  const target = req.app.locals.trainProcess?.pid + 1;
  spawnSync("kill", ["-9", `${target}`]);
  // req.app.locals.trainProcess?.kill();
  return res.json({ code: 0, msg: "停止训练任务" });
});

module.exports = router;
